MODULE optimize_eventempered
  !---------------------------------------------------------------------
  ! Completeness-optimization procedures used by Kruununhaka
  !---------------------------------------------------------------------
  ! Copyright (C) 2011 Jussi Lehtola,  University of Helsinki
  !
  ! This program is distributed under the terms of the GNU General
  ! Public License as published by the Free Software Foundation;
  ! either version 2 of the License, or (at your option) any later version.
  !
  ! This program is distributed in the hope that it will be useful,
  ! but WITHOUT ANY WARRANTY; without even the implied warranty of
  ! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. In addition, the 
  ! author must be properly cited whenever the program or some parts 
  ! originated on it are used. 
  !
  !---------------------------------------------------------------------
  USE definitions
  USE error
  USE random
  USE completeness_profile, ONLY : comp_prof
  USE measure
  USE sorting

  IMPLICIT NONE
  PRIVATE
  PUBLIC :: form_eventempered, opt_eventempered
  PUBLIC :: form_welltempered, opt_welltempered

CONTAINS
  !---------------------------------------------------------------------
  SUBROUTINE form_eventempered(alpha, beta, X)
    ! Form even-tempered set
    IMPLICIT NONE
    REAL(KIND=prec), INTENT(IN) :: alpha, beta
    REAL(KIND=prec), DIMENSION(:), INTENT(OUT) :: X

    INTEGER :: n, i
    n=SIZE(X)

    ! In an even tempered set the exponents are distributed according to
    ! zeta(i) = alpha * beta^(i-1)

    X(1)=alpha
    DO i = 2, n
       X(i)=X(i-1)*beta
    END DO
    CALL sort_desc(X)
  END SUBROUTINE form_eventempered
  !---------------------------------------------------------------------
  SUBROUTINE form_welltempered(alpha, beta, gamma, delta, X)
    ! Form well-tempered set.

    ! See e.g.
    ! S. Huzinaga and M. Klobukowski, "Well-tempered Gaussian basis
    ! sets for the calculation of matrix Hartree-Fock wavefunctions",
    ! Chem. Phys. Lett. 212, pages 260-264 (1993)

    IMPLICIT NONE
    REAL(KIND=prec), INTENT(IN) :: alpha, beta, gamma, delta
    REAL(KIND=prec), DIMENSION(:), INTENT(OUT) :: X

    INTEGER :: n, i
    n=SIZE(X)

!    ! WT-F1
!    DO i = 1, n
!       X(i)=alpha*beta**REAL(i-1,prec)*(1.0_prec + gamma*(REAL(i,prec)/REAL(n,prec))**delta)
!    END DO

    ! WT-F2, better [in energy] than WT-F1
    IF(n>=1) X(1)=alpha
    IF(n>=2) X(2)=alpha*beta
    DO i=3,N
       X(i)=X(i-1)*beta*(1.0_prec + gamma*(REAL(i,prec)/REAL(n,prec))**delta)
    END DO
    
    CALL sort_desc(X)
  END SUBROUTINE form_welltempered
  !---------------------------------------------------------------------
  FUNCTION opt_eventempered(l, optmin, optmax, X, C, usesto, measure, verbose) RESULT(mincost)
    ! Optimize the primitives using even-tempered formula
    IMPLICIT NONE
    ! Value of angular momentum
    INTEGER, INTENT(IN) :: l
    ! Exponent range to optimize in
    REAL(KIND=prec), INTENT(IN) :: optmin, optmax
    ! Even-tempered exponents to optimize
    REAL(KIND=prec), DIMENSION(:), INTENT(OUT) :: X
    ! Contraction coefficients
    REAL(KIND=prec), DIMENSION(:,:), INTENT(IN) :: C
    ! Use STOs to probe?
    LOGICAL, INTENT(IN) :: usesto
    ! Which measure to use
    INTEGER, INTENT(IN) :: measure
    ! Verbose operation?
    LOGICAL, INTENT(IN) :: verbose

    ! Probing exponents
    REAL(KIND=prec), DIMENSION(optpoints) :: alphas
    ! Alpha and beta
    REAL(KIND=prec), DIMENSION(2) :: logcurpos
    ! Step sizes
    REAL(KIND=prec), DIMENSION(2) :: stepsize
    ! New positions and values
    REAL(KIND=prec), DIMENSION(-1:1,-1:1,2) :: newpos
    REAL(KIND=prec), DIMENSION(-1:1,-1:1) :: cost
    ! When to break off due to small gradient
    REAL(KIND=prec), PARAMETER :: eps=EPSILON(stepsize)
    ! When to break off due to small change of cost
    REAL(KIND=prec), PARAMETER :: tol=1e-30
    ! RMS step size
    REAL(KIND=prec) :: rmsstep
    
    ! Current iteration
    INTEGER :: iter

    ! Time variables
    INTEGER :: time1, time2, cpucount
    ! Amount of exponents
    INTEGER :: nx

    ! Loop index
    INTEGER :: idir, jdir, j
    ! Minimum cost
    REAL(KIND=prec) :: mincost, oldcost
    ! Location of minimum
    INTEGER :: imin, jmin

    mincost=0.0_prec
    nx=SIZE(X)

    ! Fill probing exponents
    CALL get_alphas(optmin,optmax,alphas)

    ! Determine starting point.
    logcurpos(2) = LOG(optmax/optmin)/(REAL(nx,prec))
    logcurpos(1) = LOG(optmin)+0.5_prec*logcurpos(2)

    ! Initial step sizes
    stepsize=1.0_prec

    ! Fill exponents
    CALL form_eventempered(EXP(logcurpos(1)),EXP(logcurpos(2)),X)

    IF(verbose) THEN
       ! Print values
       WRITE(*,'(/,A)') 'Values for initial set:'
       ! Print out values for initial set
       CALL f_everything(l, X, C, alphas, usesto)
       WRITE(*,'(A,I3,A)') 'Optimizing a set of ', nx, &
            ' functions. This may take time.'
    END IF

    CALL system_clock(time1,cpucount)

    SELECT CASE(measure)
    CASE(1)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'Max deviation','RMS step'
       oldcost=f_maxdev(l,X,C,alphas,usesto)
       mincost=oldcost
    CASE(2)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'RMS deviation','RMS step'
       oldcost=f_lsq(l,X,C,alphas,usesto)
       mincost=oldcost
    CASE(3)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'Integral','RMS step'
       oldcost=f_sum(l,X,C,alphas,usesto)
       mincost=oldcost
    END SELECT

    iter=1
    ! Optimization loop
    DO WHILE(MAXVAL(stepsize)>eps)
       ! Increase step sizes
       stepsize=4.0_prec*stepsize

       ! Determine step size
       DO WHILE(stepsize(1)>eps .AND. stepsize(2)>eps)
          ! Determine new positions and costs
          DO idir=-1,1
             DO jdir=-1,1
                ! New position is
                newpos(idir,jdir,:)=(/logcurpos(1)+idir*stepsize(1),logcurpos(2)+jdir*stepsize(2)/)
                ! Exponents are
                CALL form_eventempered(EXP(newpos(idir,jdir,1)),EXP(newpos(idir,jdir,2)),X)
                ! and cost here is
                
                SELECT CASE(measure)
                CASE(1)
                   cost(idir,jdir)=f_maxdev(l,X,C,alphas,usesto)
                CASE(2)
                   cost(idir,jdir)=f_lsq(l,X,C,alphas,usesto)
                CASE(3)
                   cost(idir,jdir)=f_sum(l,X,C,alphas,usesto)
                END SELECT
             END DO
          END DO
          
          ! Do stepsizes need adjusting?
          IF(cost(-1,0)>cost(0,0) .AND. cost(1,0)>cost(0,0)) THEN
             stepsize(1)=stepsize(1)/2.0_prec
          ELSEIF(cost(0,-1)>cost(0,0) .AND. cost(0,1)>cost(0,0)) THEN
             stepsize(2)=stepsize(2)/2.0_prec
          ELSE
             ! Converged stepsize.
             EXIT
          END IF
       END DO

       ! Store old cost
       oldcost=mincost
       ! Get minimum cost
       mincost=cost(0,0)
       imin=0
       jmin=0
       DO idir=-1,1
          DO jdir=-1,1
             IF(cost(idir,jdir)<mincost) THEN
                imin=idir
                jmin=jdir
                mincost=cost(idir,jdir)
             END IF
          END DO
       END DO

       ! Go to location of minimum cost
       logcurpos=newpos(imin,jmin,:)

       ! Printout if necessary
       IF(verbose .AND. MOD(iter,istep_even) == 1) THEN
          ! Compute rms step size
          rmsstep=0.0_prec
          DO j=1,2
             rmsstep=rmsstep + stepsize(j)*stepsize(j)
          END DO
          rmsstep=SQRT(rmsstep/2)

          WRITE(*,'(A10,I10,E15.6,E15.6)') ' ', iter, mincost, rmsstep
       END IF
       iter=iter+1

       ! Is the change in the measure significant?
       IF(ABS(oldcost-mincost)<tol) THEN
          EXIT
       END IF
    END DO

    CALL system_clock(time2,cpucount)

    IF(verbose) THEN
       WRITE (*,'(A,E11.4,A,I4,A,/)') 'Converged to ', mincost, ' in ',iter,' iterations.'
       WRITE (*,'(A,2E11.4,/)') 'Parameters are ',EXP(logcurpos)
       WRITE (*,'(A,F5.1,A,/)') 'Optimization took',&
            (time2-time1)*1.0/cpucount,' seconds.'
    END IF
    ! Form optimal set
    CALL form_eventempered(EXP(logcurpos(1)),EXP(logcurpos(2)),X)
    ! Sort it in descending order
    CALL sort_desc(X)

    ! Print values
    IF(verbose) THEN
       WRITE(*,'(A)') 'Values for optimized set:'
       ! Compute measures of goodness for optimal set
       CALL f_everything(l, X, C, alphas, usesto)
    END IF
  END FUNCTION opt_eventempered
  !---------------------------------------------------------------------
  FUNCTION opt_welltempered(l, optmin, optmax, X, C, usesto, measure, verbose) RESULT(curiter)
    ! Optimize the primitives by the well-tempered formula
    IMPLICIT NONE
    ! Value of angular momentum
    INTEGER, INTENT(IN) :: l
    ! Exponent range to optimize in
    REAL(KIND=prec), INTENT(IN) :: optmin, optmax
    ! Exponents to optimize
    REAL(KIND=prec), DIMENSION(:), INTENT(INOUT) :: X
    ! Contraction coefficients
    REAL(KIND=prec), DIMENSION(:,:), INTENT(IN) :: C
    ! Use STOs to probe?
    LOGICAL, INTENT(IN) :: usesto
    ! Which measure to use
    INTEGER, INTENT(IN) :: measure
    ! Verbose operation?
    LOGICAL, INTENT(IN) :: verbose

    ! Probing exponents
    REAL(KIND=prec), DIMENSION(optpoints) :: alphas

    ! All parameters have to be posive => optimization in logarithmic
    ! space
    ! Number of parameters
    INTEGER, PARAMETER :: npar=4
    ! Current parameters
    REAL(KIND=prec), DIMENSION(npar) :: params

    ! Logarithms of parameters
    REAL(KIND=prec), DIMENSION(npar) :: logparams
    ! Step sizes
    REAL(KIND=prec), DIMENSION(npar) :: stepsize
    ! RMS step
    REAL(KIND=prec) :: rmsstep
    ! When to break off due to small gradient
    REAL(KIND=prec), PARAMETER :: eps=EPSILON(stepsize)
    ! When to break off due to small total change of cost
    REAL(KIND=prec), PARAMETER :: tol=1e-30
    ! Current iteration
    INTEGER :: iter

    ! Time variables
    INTEGER :: time1, time2, cpucount
    ! Amount of exponents
    INTEGER :: nx

    ! Loop indices
    INTEGER :: ip, j
    ! Trial parameters
    REAL(KIND=prec), DIMENSION(npar,-1:1) :: partrial
    ! Trial exponents
    REAL(KIND=prec), DIMENSION(SIZE(X),-1:1) :: Xtrial

    ! Costs of current and last iteration over all exponents
    REAL(KIND=prec) :: curiter, lastiter
    ! Current cost
    REAL(KIND=prec) :: curcost=0.0_prec
    ! Costs of varying current parameter
    REAL(KIND=prec), DIMENSION(-1:1) :: trialcost

    ! Position and value of minimum
    INTEGER :: minpos
    REAL(KIND=prec) :: minval

    curiter=0.0_prec
    nx=SIZE(X)

    ! Fill probing exponents
    CALL get_alphas(optmin,optmax,alphas)
    ! Print values
    IF(verbose) THEN
       WRITE(*,'(/,A)') 'Values for initial set:'
       ! Print out values for initial set
       CALL f_everything(l, X, C, alphas, usesto)
       WRITE(*,'(A,I3,A)') 'Optimizing a set of ', nx, &
            ' functions. This may take time.'
    END IF
       
    CALL system_clock(time1,cpucount)

    ! Initialize alpha and beta parameters
    params(2) = EXP(LOG(optmax/optmin)/REAL(nx,prec))
    params(1) = optmin*SQRT(params(2))
    params(3) = 1.0_prec
    params(4) = 1.0_prec

    logparams=LOG(params)

    ! Initialize exponents
    CALL form_welltempered(params(1),params(2),params(3),params(4),X)

    ! Initial step sizes
    stepsize=1.0_prec;

    ! Calculate initial measure of goodness
    SELECT CASE(measure)
    CASE(1)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'Max deviation','RMS step'
       curiter=f_maxdev(l,X,C,alphas,usesto)
       curcost=curiter
    CASE(2)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'RMS deviation','RMS step'
       curiter=f_lsq(l,X,C,alphas,usesto)
       curcost=curiter
    CASE(3)
       IF(verbose) WRITE(*,'(/,A10,A10,A15,A15)') ' ', 'Iteration', 'Integral','RMS step'
       curiter=f_sum(l,X,C,alphas,usesto)
       curcost=curiter
    END SELECT

    iter=1
    ! Optimization loop
    DO WHILE(MAXVAL(stepsize)>eps)
       ! Increase step sizes
       stepsize=4.0_prec*stepsize

       ! Store old value of cost function
       lastiter=curiter

       ! Loop over parameters
       DO ip=1,npar
          ! Current exponents
          Xtrial(:,0)=X
          partrial(:,0)=logparams

          ! Determine step size for current parameter
          DO WHILE(stepsize(ip)>eps)
             ! Calculate trials
             DO j=-1,1,2
                partrial(:,j)=logparams
                partrial(ip,j)=partrial(ip,j)+j*stepsize(ip)
                CALL form_welltempered(EXP(partrial(1,j)),EXP(partrial(2,j)),EXP(partrial(3,j)),EXP(partrial(4,j)),Xtrial(:,j))
             END DO

             ! and costs
             SELECT CASE(measure)
             CASE(1)
                DO j=-1,1,2
                   trialcost(j)=f_maxdev(l,Xtrial(:,j),C,alphas,usesto)
                END DO
                trialcost(0)=curcost
             CASE(2)
                DO j=-1,1,2
                   trialcost(j)=f_lsq(l,Xtrial(:,j),C,alphas,usesto)
                END DO
                trialcost(0)=curcost
             CASE(3)
                DO j=-1,1,2
                   trialcost(j)=f_sum(l,Xtrial(:,j),C,alphas,usesto)
                END DO
                trialcost(0)=curcost
             END SELECT

             ! Find out minimum position
             minval=trialcost(0)
             minpos=0
             DO j=-1,1,2
                IF(trialcost(j)<minval) THEN
                   minpos=j
                   minval=trialcost(j)
                END IF
             END DO

             ! Does stepsize need adjusting?
             IF(minpos==0) THEN
                stepsize(ip)=stepsize(ip)/2.0_prec
             ELSE
                ! No, break while loop
                EXIT
             END IF
          END DO

          ! If all went well, we have a lower value
          IF(trialcost(-1)<curcost) THEN
             curcost=trialcost(-1)
             X=Xtrial(:,-1)
             logparams=partrial(:,-1)
             params=EXP(partrial(:,-1))
          ELSE IF(trialcost(1)<curcost) THEN
             ! Lower value is on the right
             curcost=trialcost(1)
             X=Xtrial(:,1)
             logparams=partrial(:,1)
             params=EXP(partrial(:,1))
!          ELSE
!             WRITE (*,*) 'Parameter ',ip,' is stagnated.'
          END IF
       END DO

       ! Store current value of cost
       curiter=curcost
       
       ! Printout if necessary
       IF(verbose .AND. MOD(iter,istep) == 1) THEN
          ! Compute rms step size
          rmsstep=0.0_prec
          DO j=1,npar
             rmsstep=rmsstep + stepsize(j)*stepsize(j)
          END DO
          rmsstep=SQRT(rmsstep/npar)

          WRITE(*,'(A10,I10,E15.6,E15.6)') ' ', iter, curiter, rmsstep
       END IF
       iter=iter+1

       ! Is the change in the measure significant?
       IF(ABS(curiter-lastiter)<tol) THEN
          EXIT
       END IF
    END DO

    CALL system_clock(time2,cpucount)
    
    IF(verbose) THEN
       WRITE (*,'(A,E11.4,A,I4,A,/)') 'Converged to ', curiter, ' in ',iter,' iterations.'
       WRITE (*,'(A,4E11.4,/)') 'Parameters are ',params
       WRITE (*,'(A,F5.1,A,/)') 'Optimization took',&
            (time2-time1)*1.0/cpucount,' seconds.'
    END IF

    ! Sort optimal set in descending order
    CALL sort_desc(X)

    ! Print values
    IF(verbose) THEN
       WRITE(*,'(A)') 'Values for optimized set:'
       ! Compute measures of goodness for optimal set
       CALL f_everything(l, X, C, alphas, usesto)
    END IF
  END FUNCTION opt_welltempered
END MODULE optimize_eventempered
